import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def create_transformation_matrix(rotation, translation):
    """
    创建4x4的刚体变换矩阵
    参数:
        rotation: 3x3的旋转矩阵
        translation: 3x1的平移向量
    返回:
        4x4的变换矩阵
    """
    transformation = np.eye(4)
    transformation[:3, :3] = rotation
    transformation[:3, 3] = translation
    return transformation

def get_transformation_matrix(transform):
    """
    将12维的刚体矩阵转换为4x3的矩阵
    参数:
        transform: 12维的numpy数组，表示刚体变换矩阵的前3行
    返回:
        3x4的变换矩阵
    """
    return transform.reshape((3, 4))

def draw_coordinate_axes(ax, transformation, scale=1.0, label=""):
    """
    在3D坐标系中绘制坐标轴
    参数:
        ax: matplotlib的3D坐标轴对象
        transformation: 3x4的变换矩阵
        scale: 坐标轴的长度缩放因子
        label: 坐标系的标签
    """
    # 提取旋转和平移
    rotation = transformation[:, :3]
    translation = transformation[:, 3]
    
    # 计算坐标轴的终点
    x_axis = translation + scale * rotation[:, 0]
    y_axis = translation + scale * rotation[:, 1]
    z_axis = translation + scale * rotation[:, 2]
    
    # 绘制坐标轴
    ax.plot([translation[0], x_axis[0]], [translation[1], x_axis[1]], [translation[2], z_axis[2]], 'r-', label=f"{label} X轴")
    ax.plot([translation[0], y_axis[0]], [translation[1], y_axis[1]], [translation[2], y_axis[2]], 'g-', label=f"{label} Y轴")
    ax.plot([translation[0], z_axis[0]], [translation[1], z_axis[1]], [translation[2], z_axis[2]], 'b-', label=f"{label} Z轴")
    
    # 标记坐标系中心
    ax.scatter(translation[0], translation[1], translation[2], color='black', label=f"{label} 原点")
    
    # 添加箭头
    ax.quiver(translation[0], translation[1], translation[2], 
              rotation[0, 0], rotation[1, 0], rotation[2, 0], 
              color='r', length=scale*0.9, arrow_length_ratio=0.1)
    ax.quiver(translation[0], translation[1], translation[2], 
              rotation[0, 1], rotation[1, 1], rotation[2, 1], 
              color='g', length=scale*0.9, arrow_length_ratio=0.1)
    ax.quiver(translation[0], translation[1], translation[2], 
              rotation[0, 2], rotation[1, 2], rotation[2, 2], 
              color='b', length=scale*0.9, arrow_length_ratio=0.1)

def visualize_transformations(transform1, transform2):
    """
    可视化两个刚体变换矩阵
    参数:
        transform1: 第一个刚体变换矩阵（12维numpy数组）
        transform2: 第二个刚体变换矩阵（12维numpy数组）
    """
    # 将输入转换为3x4矩阵
    t1 = get_transformation_matrix(transform1)
    t2 = get_transformation_matrix(transform2)
    
    # 创建3D坐标轴
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 设置坐标轴范围和标签
    ax.set_xlabel('X (米)')
    ax.set_ylabel('Y (米)')
    ax.set_zlabel('Z (米)')
    ax.set_title('3D刚体变换矩阵可视化')
    
    # 设置坐标轴范围 (可根据变换矩阵的平移部分自动调整)
    min_val = min(t1[:, 3].min(), t2[:, 3].min()) - 1
    max_val = max(t1[:, 3].max(), t2[:, 3].max()) + 1
    ax.set_xlim([min_val, max_val])
    ax.set_ylim([min_val, max_val])
    ax.set_zlim([min_val, max_val])
    
    # 绘制两个坐标系
    draw_coordinate_axes(ax, t1, scale=0.5, label="变换1")
    draw_coordinate_axes(ax, t2, scale=0.5, label="变换2")
    
    # 添加图例和网格
    ax.legend()
    ax.grid(True)
    
    # 显示图形
    plt.show()

# 示例输入
# 创建两个示例变换矩阵
# 变换1：恒等变换，原点在(0,0,0)
transform1 = np.array([1, 0, 0, 0,
                       0, 1, 0, 0,
                       0, 0, 1, 0])

# 变换2：旋转+平移
rotation2 = np.array([[0, 0, 1],
                      [1, 0, 0],
                      [0, 1, 0]])
translation2 = np.array([2, 3, 1])
transform2 = np.zeros(12)
transform2[:9] = rotation2.flatten()
transform2[9:] = translation2

# 可视化变换
visualize_transformations(transform1, transform2)